#pragma once

#include <algorithm>
#include <numeric>
#include <vector>

#include "whimed.h"

// TODO: citation

/**
 * @brief Implementations based on R package robustbase
 *
 * @cite Maechler M, Rousseeuw P, Croux C, Todorov V, Ruckstuhl A, Salibian-Barrera M, Verbeke T, Koller M, Conceicao EL, Anna di Palma M (2021).
 * robustbase: Basic Robust Statistics. R package version 0.93-8, http://robustbase.r-forge.r-project.org/
 *
 */

namespace mdtx::online::algo
{
    template <typename T>
    class sn
    {
    protected:
        ptrdiff_t n;
        std::vector<T> a2;
        static constexpr double cn[10] = {0.000, 0.000, 0.743, 1.851, 0.954,
                                          1.351, 0.993, 1.198, 1.005, 1.131};

    public:
        explicit sn(ptrdiff_t n)
            : n(n),
              a2(n)
        {
        }

        T get_cn(ptrdiff_t n)
        {
            if (n < 10)
            {
                return T(cn[n]);
            }
            else
            {
                if (n % 2)
                {
                    return n / T(n - 0.9);
                }
                else
                {
                    return T(1);
                }
            }
        }

        template <typename iter>
        auto run(iter it)
        {
            T medA, medB;
            ptrdiff_t i, diff, half, Amin, Amax, even, length;
            ptrdiff_t leftA, leftB, nA, nB, tryA, tryB, rightA, rightB;
            ptrdiff_t n1_2;

            auto get = [it](ptrdiff_t idx)
            {
                return *(it + idx);
            };

            a2[0] = get(n / 2) - get(0);
            n1_2 = (n + 1) / 2;

            // first half for() loop:
            for (i = 2; i <= n1_2; ++i)
            {
                nA = i - 1;
                nB = n - i;
                diff = nB - nA;
                leftA = leftB = 1;
                rightA = rightB = nB;
                Amin = diff / 2 + 1;
                Amax = diff / 2 + nA;

                while (leftA < rightA)
                {
                    length = rightA - leftA + 1;
                    even = 1 - length % 2;
                    half = (length - 1) / 2;
                    tryA = leftA + half;
                    tryB = leftB + half;
                    if (tryA < Amin)
                    {
                        rightB = tryB;
                        leftA = tryA + even;
                    }
                    else
                    {
                        if (tryA > Amax)
                        {
                            rightA = tryA;
                            leftB = tryB + even;
                        }
                        else
                        {
                            medA = get(i - 1) - get(i - tryA + Amin - 2);
                            medB = get(tryB + i - 1) - get(i - 1);
                            if (medA >= medB)
                            {
                                rightA = tryA;
                                leftB = tryB + even;
                            }
                            else
                            {
                                rightB = tryB;
                                leftA = tryA + even;
                            }
                        }
                    }
                } // while

                if (leftA > Amax)
                {
                    a2[i - 1] = get(leftB + i - 1) - get(i - 1);
                }
                else
                {
                    medA = get(i - 1) - get(i - leftA + Amin - 2);
                    medB = get(leftB + i - 1) - get(i - 1);
                    a2[i - 1] = std::min(medA, medB);
                }
            }

            // second half for() loop:
            for (i = n1_2 + 1; i <= n - 1; ++i)
            {
                nA = n - i;
                nB = i - 1;
                diff = nB - nA;
                leftA = leftB = 1;
                rightA = rightB = nB;
                Amin = diff / 2 + 1;
                Amax = diff / 2 + nA;

                while (leftA < rightA)
                {
                    length = rightA - leftA + 1;
                    even = 1 - length % 2;
                    half = (length - 1) / 2;
                    tryA = leftA + half;
                    tryB = leftB + half;
                    if (tryA < Amin)
                    {
                        rightB = tryB;
                        leftA = tryA + even;
                    }
                    else
                    {
                        if (tryA > Amax)
                        {
                            rightA = tryA;
                            leftB = tryB + even;
                        }
                        else
                        {
                            medA = get(i + tryA - Amin) - get(i - 1);
                            medB = get(i - 1) - get(i - tryB - 1);
                            if (medA >= medB)
                            {
                                rightA = tryA;
                                leftB = tryB + even;
                            }
                            else
                            {
                                rightB = tryB;
                                leftA = tryA + even;
                            }
                        }
                    }
                } // while

                if (leftA > Amax)
                {
                    a2[i - 1] = get(i - 1) - get(i - leftB - 1);
                }
                else
                {
                    medA = get(i + leftA - Amin) - get(i - 1);
                    medB = get(i - 1) - get(i - leftB - 1);
                    a2[i - 1] = std::min(medA, medB);
                }
            }
            a2[n - 1] = get(n - 1) - get(n1_2 - 1);
            std::nth_element(a2.begin(), a2.begin() + n1_2 - 1, a2.end());

            return a2[n1_2 - 1];
        }
    };

    template <typename T>
    class qn
    {
    public:
        using tvect = std::vector<T>;
        using ivect = std::vector<ptrdiff_t>;
        using tpair = std::pair<T, ptrdiff_t>;
        using pvect = std::vector<tpair>;

    protected:
        ptrdiff_t n;
        pvect pwork;
        tvect work;
        ivect left, right, p, q;
        static constexpr double dn[10] = {0.000, 0.000, 0.399, 0.994, 0.512,
                                          0.844, 0.611, 0.857, 0.669, 0.872};

    public:
        explicit qn(ptrdiff_t n)
            : n(n),
              pwork(n), work(n),
              left(n), right(n), p(n), q(n)
        {
        }

        T get_dn(ptrdiff_t n)
        {
            if (n < 10)
            {
                return T(dn[n]);
            }
            else
            {
                if (n % 2)
                {
                    return n / T(n + 1.4);
                }
                else
                {
                    return n / T(n + 3.8);
                }
            }
        }

        template <typename iter>
        auto run(iter it)
        {
            T trial;
            ptrdiff_t h, i, j, jj, jh;
            ptrdiff_t k, knew, nl, nr;

            auto get = [it](ptrdiff_t idx)
            {
                return *(it + idx);
            };

            h = n / 2 + 1;
            k = h * (h - 1) / 2;
            for (i = 0; i < n; ++i)
            {
                left[i] = n - i + 1;
                // the n - (i - h) is from the paper; original code had `n'
                right[i] = (i <= h) ? n : n - (i - h);
            }
            nl = n * (n + 1) / 2;
            nr = n * n;
            knew = k + nl;

            // L200:
            while (nr - nl > n)
            {
                j = 0;
                // Truncation to float: try to make sure that the same values are got later (guard bits !)
                // Jun: there is no arithmic op done on work[], and trial during whimed_i, then why would we want to truncate work[]?
                for (i = 1; i < n; ++i)
                {
                    if (left[i] <= right[i])
                    {
                        pwork[j].second = right[i] - left[i] + 1;
                        jh = left[i] + pwork[j].second / 2;
                        pwork[j].first = get(i) - get(n - jh);
                        ++j;
                    }
                }
                trial = whimed::run(pwork.begin(), j);

                j = 0;
                for (i = n - 1; i >= 0; --i)
                {
                    while (j < n && get(i) - get(n - j - 1) < trial)
                    {
                        ++j;
                    }
                    p[i] = j;
                }
                j = n + 1;
                for (i = 0; i < n; ++i)
                {
                    while (get(i) - get(n - j + 1) > trial)
                    {
                        --j;
                    }
                    q[i] = j;
                }

                auto sump = std::accumulate(p.begin(), p.end(), ptrdiff_t{});
                auto sumq = std::accumulate(q.begin(), q.end(), ptrdiff_t{}) - n;

                if (knew <= sump)
                {
                    std::copy(p.begin(), p.end(), right.begin());
                    nr = sump;
                }
                else if (knew > sumq)
                {
                    std::copy(q.begin(), q.end(), left.begin());
                    nl = sumq;
                }
                else
                {
                    // sump < knew <= sumq
                    return trial;
                }
            } // while

            j = 0;
            for (i = 1; i < n; ++i)
            {
                for (jj = left[i]; jj <= right[i]; ++jj)
                {
                    work[j] = get(i) - get(n - jj);
                    // j will be = sum_{i=2}^n (right[i] - left[i] + 1)_{+}
                    j++;
                }
            }

            knew -= (nl + 1);
            std::nth_element(work.begin(), work.begin() + knew, work.begin() + j);
            return work[knew];
        }
    };
}
